#include "stdafx.h"
#include "ThreadIOCP.h"
#include <iostream>
#include "IOCompletionPort.h"
#include "CircleDequeueBase.h"
#include "DataDefinition.h"
#include "DateTime.h"
#include "../DzhNewTQSDlg/ConfigInfoInterface.h"
#include <string>

const char ARRAY_USER_ID[]= "715jdr";
const char ARRAY_PASSWORD[] = "$%^wgdzh";
const char ARRAY_SUCCUSSFUL_MSG[] = "userpasssuce";
const unsigned int USER_ID_LENGTH = 6;
const unsigned int USER_PASSWORD_LENGTH = 8;
const unsigned int SUCCESSFUL_MESSAGE_LENGTH = 12;

ThreadIOCP::ThreadIOCP(map<int,CircleDequeueBase*> p_mapCircleQueue, ConnectionInformation* p_pcClientList)
	:m_pcClientList(p_pcClientList)
	,m_mapCircleQueue(p_mapCircleQueue)
	,m_lstWaitingClients()
{
}

bool ThreadIOCP::VerifyUser(char* p_szUserInfo, unsigned int p_uiLength)
{
	// Verify parameters
	if(NULL == p_szUserInfo)
	{
		LOG(ERROR) << "Parameter p_szUserInfo is NULL";
		return false;
	}
	if(p_uiLength < USER_ID_LENGTH + USER_PASSWORD_LENGTH)
	{
		LOG(WARNING) << "User verification is incorrect. Message length(" << p_uiLength << ") is less than configured";
		return false;
	}

	// Verify User ID and Password
	if(0 != memcmp(p_szUserInfo,ARRAY_USER_ID,USER_ID_LENGTH))
	{
		LOG(WARNING) << "Cannot pass authentication. Incorrect user ID or password";
		return false;
	}

	char* pPassword = (p_szUserInfo + 9);
	for(int i=0; i < USER_PASSWORD_LENGTH; i++)
	{
		if(ARRAY_PASSWORD[i] != (char)((int)(pPassword[i])+13+i))
		{
			LOG(WARNING) << "Cannot pass authentication";
			return false;
		}
	}

	return true;
}

char* ThreadIOCP::GetBufferFromMemoryPool(unsigned int p_uilength, unsigned int p_uiIndex)
{
	MemoryPool* pcMemoryPool = MemoryPool::Instance();

	// Verify parameters
	if(NULL == pcMemoryPool)
	{
		LOG(ERROR) << "Instance of memory pool is NULL";
		return NULL;
	}
	if(0 == p_uilength)
	{
		LOG(ERROR) << "Buffer length is 0";
		return NULL;
	}

	// Allocate buffer from memory pool
	return pcMemoryPool->Allocate(p_uilength, p_uiIndex);
}

PER_IO_CONTEXT* ThreadIOCP::GetNewIOConetext(char* p_szBuffer, unsigned int p_uiBufferLength, OPERATION_TYPE p_eType, SOCKET p_stSocket, bool p_bNeedRelease)
{
	// Verify parameters
	if(NULL == p_szBuffer)
	{
		LOG(ERROR) << "Parameter p_szBuffer is NULL";
		return NULL;
	}
	if(0 == p_uiBufferLength)
	{
		LOG(ERROR) << "Parameter p_uiBufferLength is 0";
		return NULL;
	}

	// New PER_IO_CONTEXT and assign value
	PER_IO_CONTEXT* pstIOContext = new PER_IO_CONTEXT();
	if(NULL == pstIOContext)
	{
		LOG(ERROR) << "Allocate IO context fail";
	}

	pstIOContext->m_wsaBuf.buf = p_szBuffer;
	pstIOContext->m_wsaBuf.len = p_uiBufferLength;
	pstIOContext->m_OpType = p_eType;
	pstIOContext->m_sockAccept = p_stSocket;
	pstIOContext->m_bNeedRelease = p_bNeedRelease;

	return pstIOContext;
}

bool ThreadIOCP::DoAccept(PER_SOCKET_CONTEXT* p_pstSocketContext, PER_IO_CONTEXT* p_pstIoContext)
{
	IOCompletionPort* pInstance = IOCompletionPort::Instance();

	// Verify user ID and password is correct, if not correct close the socket
	if(false == this->VerifyUser(p_pstIoContext->m_wsaBuf.buf, p_pstIoContext->m_wsaBuf.len))
	{
		//LOG(WARNING) << "Client " << inet_ntoa(p_pstSocketContext->m_ClientAddr.sin_addr) << "verified failed";
		this->ReAccept(p_pstIoContext);
		return false;
	}

	// Associate with ICOP and add it to client list in IOCP
	string strClientIP;
	unsigned int uiClientPort = 0;
	if(false == pInstance->AssociateWithIOCP(p_pstIoContext,strClientIP,uiClientPort))
	{
		LOG(ERROR) << "AssociateWithIOCP() Error : " << WSAGetLastError();
		this->ReAccept(p_pstIoContext);
		return false;
	}

	// Set client ip and port to client information
	m_pcClientList->SetRemoteIP(p_pstIoContext->m_sockAccept,strClientIP);
	m_pcClientList->SetRemotePort(p_pstIoContext->m_sockAccept,uiClientPort);
	m_pcClientList->SetLocalIP(p_pstIoContext->m_sockAccept,CConfigInfoInterface::Instance()->GetServerSetting().m_strHostIP);
	m_pcClientList->SetLocalPort(p_pstIoContext->m_sockAccept,CConfigInfoInterface::Instance()->GetServerSetting().m_uiHostPort);
	LOG(INFO) << "Client " << strClientIP << ":" << uiClientPort << " connected";

	// WSASend successful message
	char* pszFeedback = this->GetBufferFromMemoryPool(SUCCESSFUL_MESSAGE_LENGTH, AUTH_MESSAGE_ID);
	if(NULL == pszFeedback)
	{
		LOG(ERROR) << "Cannot allocate memory from memory pool";
		this->ReAccept(p_pstIoContext);
		return false;
	}
					
	// Copy message to buffer
	memcpy(pszFeedback,ARRAY_SUCCUSSFUL_MSG,SUCCESSFUL_MESSAGE_LENGTH);
	PER_IO_CONTEXT* pstSendContext = this->GetNewIOConetext(pszFeedback,SUCCESSFUL_MESSAGE_LENGTH,AUTH_POSTED,p_pstIoContext->m_sockAccept,true);
	if(false == pInstance->PostSend(pstSendContext))
	{
		LOG(ERROR) << "Post send error. Socket : " << pstSendContext->m_sockAccept;
		
		// Release related resource
		pInstance->RemoveContext(pstSendContext->m_sockAccept);
		MemoryPool::Instance()->Release(pstSendContext->m_wsaBuf.buf);
		m_pcClientList->DeleteClient(pstSendContext->m_sockAccept);
		RELEASE(pstSendContext);

		// Continue to accept
		p_pstIoContext->ResetBuffer();
		if(false == IOCompletionPort::Instance()->PostAccept(p_pstIoContext))
		{
			LOG(ERROR) << "Post accept error";
		}

		return false;
	}

	// Continue to listen and post WSAReceive
	char* pszReceiveBuffer = (char*)malloc(sizeof(FileType));
	if(NULL == pszReceiveBuffer)
	{
		// Close socket
		LOG(ERROR) << "Cannot allocate memory from memory pool";
		this->ReAccept(p_pstIoContext);
		return false;
	}

	PER_IO_CONTEXT* pstReceiveContext = this->GetNewIOConetext(pszReceiveBuffer,sizeof(FileType),RECV_POSTED,p_pstIoContext->m_sockAccept,true);
	if(false == pInstance->PostRecv(pstReceiveContext))
	{
		LOG(ERROR) << "Post receive error";
		pInstance->RemoveContext(pstReceiveContext->m_sockAccept);
		MemoryPool::Instance()->Release(pstReceiveContext->m_wsaBuf.buf);
		m_pcClientList->DeleteClient(pstReceiveContext->m_sockAccept);
		RELEASE(pstReceiveContext);

		// Continue to accept
		p_pstIoContext->ResetBuffer();
		if(false == IOCompletionPort::Instance()->PostAccept(p_pstIoContext))
		{
			LOG(ERROR) << "Post accept error";
		}
		return false;
	}

	// Continue to accept
	p_pstIoContext->ResetBuffer();
	if(false == pInstance->PostAccept(p_pstIoContext))
	{
		LOG(ERROR) << "Post accept error";
		return false;
	}

	return true;
}

bool ThreadIOCP::ReAccept(PER_IO_CONTEXT* p_pstIoContext)
{
	// Close connection
	RELEASE_SOCKET(p_pstIoContext->m_sockAccept);

	// Continue to accept
	p_pstIoContext->ResetBuffer();
	if(false == IOCompletionPort::Instance()->PostAccept(p_pstIoContext))
	{
		LOG(ERROR) << "Post accept error";
		return false;
	}

	return true;
}

bool ThreadIOCP::DoReceive(PER_SOCKET_CONTEXT* p_pstSocketContext, PER_IO_CONTEXT* p_pstIoContext)
{
	// Verify parameters
	if(NULL == p_pstSocketContext)
	{
		LOG(ERROR) << "Parameter p_pstSocketContext is NULL";
		return false;
	}
	if(NULL == p_pstIoContext)
	{
		LOG(ERROR) << "Parameter p_pstIoContext is NULL";
		return false;
	}
	
	IOCompletionPort* pInstance = IOCompletionPort::Instance();

	if (sizeof(DataHead) == p_pstIoContext->m_wsaBuf.len)
	{
		// Re-active heart beat
		DateTime clsDateTime;
		clsDateTime.SetToCurrentDateTime();
		m_pcClientList->SetTime(p_pstIoContext->m_sockAccept,clsDateTime);
		free(p_pstIoContext->m_wsaBuf.buf);
		RELEASE(p_pstIoContext)

		// Add waiting clients to relevant structure
		while(m_lstWaitingClients.size() != 0)
		{
			if (MemoryPool::Instance()->GetMemoryStatus() > 0.7)
			{
				LOG(WARNING) << "Memory pool is full . Waiting for next loop. Size : " <<  m_lstWaitingClients.size();
				break;
			} 
			else
			{
				SOCKET iClientSocket = m_lstWaitingClients.front();
				m_lstWaitingClients.pop_front();
				FileType stWaitingFileType = m_pcClientList->GetFileType(iClientSocket);
				this->DecodeFileType(iClientSocket,stWaitingFileType);
				LOG(WARNING) << "[BeatHeart] Deal with left clients. Left " << m_lstWaitingClients.size() << " clients.";
			}
		}

		return true;
	}

	// Receive file list and compress type request
	if (sizeof(FileType) != p_pstIoContext->m_wsaBuf.len)
	{
		LOG(ERROR) << "File type length dismatch";
		m_pcClientList->DeleteClient(p_pstIoContext->m_sockAccept);
		pInstance->RemoveContext(p_pstIoContext->m_sockAccept);
		free(p_pstIoContext->m_wsaBuf.buf);
		RELEASE_SOCKET(p_pstIoContext->m_sockAccept);
		RELEASE(p_pstIoContext)
		return true;
	}
	else
	{
		FileType stFileType;
		memcpy(&stFileType,p_pstIoContext->m_wsaBuf.buf,sizeof(FileType));
		m_pcClientList->SetFileType(p_pstIoContext->m_sockAccept,stFileType);
		if (MemoryPool::Instance()->GetMemoryStatus() > 0.7)
		{
			LOG(WARNING) << "Memory pool is already full. Push client to waiting queue. Size : " <<  m_lstWaitingClients.size();
			m_lstWaitingClients.push_back(p_pstIoContext->m_sockAccept);
		} 
		else
		{
			// Add client information and request to relevant structure
			this->DecodeFileType(p_pstIoContext->m_sockAccept,stFileType);

			// Add waiting clients to relevant structure
			while(m_lstWaitingClients.size() != 0)
			{
				if (MemoryPool::Instance()->GetMemoryStatus() > 0.9)
				{
					LOG(WARNING) << "Memory pool is full. Waiting for next loop. Size : " <<  m_lstWaitingClients.size();
					break;
				} 
				else
				{
					SOCKET iClientSocket = m_lstWaitingClients.front();
					m_lstWaitingClients.pop_front();
					FileType stWaitingFileType = m_pcClientList->GetFileType(iClientSocket);
					this->DecodeFileType(iClientSocket,stWaitingFileType);
					LOG(WARNING) << "Deal with left clients. Left " << m_lstWaitingClients.size() << " clients.";
				}
			}
		}
	}

	free(p_pstIoContext->m_wsaBuf.buf);
	RELEASE(p_pstIoContext)

	return true;
}

void ThreadIOCP::AddClientToQueue(unsigned char p_ucFileTypeFlag,const int p_iFileType, unsigned char p_ucCompressType,SOCKET p_iClient)
{
	if((1 == p_ucFileTypeFlag) && (1 == p_ucCompressType))
	{
		map<int,CircleDequeueBase*>::iterator it = m_mapCircleQueue.find(p_iFileType);
		if(it != m_mapCircleQueue.end())
		{
			bool bRet = m_mapCircleQueue[p_iFileType]->PostData(p_iClient);
			if (false == bRet)
			{
				LOG(ERROR) << "Post data error";
				m_pcClientList->DeleteClient(p_iClient);
				IOCompletionPort::Instance()->RemoveContext(p_iClient);
				RELEASE_SOCKET(p_iClient);
				this->RemoveFromQueue(p_iClient);
			}
		}
		else
		{
			LOG(WARNING) << "No supported file request. File type : " << p_iFileType;
		}
	}
}

void ThreadIOCP::DecodeFileType(SOCKET p_iSocket, FileType p_stFileType)
{
	// Decode file type for client
	this->AddClientToQueue(p_stFileType.m_cShow2003Flag,FILETYPE_SHOW,p_stFileType.m_cAttrib,p_iSocket);
	this->AddClientToQueue(p_stFileType.m_cSjshqFlag,FILETYPE_SJSHQ,p_stFileType.m_cAttrib,p_iSocket);
	this->AddClientToQueue(p_stFileType.m_cGzlxFlag,FILETYPE_GZLX,p_stFileType.m_cAttrib,p_iSocket);
	this->AddClientToQueue(p_stFileType.m_cStbhqFlag,FILETYPE_STBHQ,p_stFileType.m_cAttrib,p_iSocket);
	this->AddClientToQueue(p_stFileType.m_c50EtfFlag,FILETYPE_50ETF,p_stFileType.m_cAttrib,p_iSocket);
	this->AddClientToQueue(p_stFileType.m_cSjsxxFlag,FILETYPE_SJSXX,p_stFileType.m_cAttrib,p_iSocket);

	this->AddClientToQueue(p_stFileType.m_cZhbFlag,FILETYPE_ZHB,p_stFileType.m_cAttrib,p_iSocket);
	this->AddClientToQueue(p_stFileType.m_cZhggqzFlag,FILETYPE_ZHGGQZ,p_stFileType.m_cAttrib,p_iSocket);
	this->AddClientToQueue(p_stFileType.m_cZhjbxxFlag,FILETYPE_ZHJBXX,p_stFileType.m_cAttrib,p_iSocket);
	this->AddClientToQueue(p_stFileType.m_cHqzFlag,FILETYPE_HQZ,p_stFileType.m_cAttrib,p_iSocket);
	this->AddClientToQueue(p_stFileType.m_cHtindexFlag,FILETYPE_HTINDEX,p_stFileType.m_cAttrib,p_iSocket);
	this->AddClientToQueue(p_stFileType.m_cHtidxconFlag,FILETYPE_HTIDXCON,p_stFileType.m_cAttrib,p_iSocket);
	this->AddClientToQueue(p_stFileType.m_cZxzsFlag,FILETYPE_ZXZS,p_stFileType.m_cAttrib,p_iSocket);
	this->AddClientToQueue(p_stFileType.m_bitHghqFlag,FILETYPE_HGHQ,p_stFileType.m_cAttrib,p_iSocket);
	this->AddClientToQueue(p_stFileType.m_bitHgxxnFlag,FILETYPE_HGXXN,p_stFileType.m_cAttrib,p_iSocket);
	this->AddClientToQueue(p_stFileType.m_bitCpxxFlag,FILETYPE_CPXX,p_stFileType.m_cAttrib,p_iSocket);
}

void ThreadIOCP::Run()
{
	IOCompletionPort* pIOCP = IOCompletionPort::Instance();

	while(1)
	{
		PER_IO_CONTEXT* pIoContext = NULL;
		PER_SOCKET_CONTEXT* pSocketContext = NULL;
		char* szMessage = NULL;

		if(false == pIOCP->GetQueuedStatus(&pSocketContext,&pIoContext))
		{
			if(IOCP_CLIENT_DISCONNECT == pIOCP->GetLastError())
			{// Client disconnect
				RemoveFromQueue(pIoContext->m_sockAccept);
				this->RemoveClient(pIoContext);
			}
			else
			{// Other error
				LOG(ERROR) << "GetQueuedStatus error:" << pIOCP->GetLastError();
				RemoveFromQueue(pIoContext->m_sockAccept);
				this->RemoveClient(pIoContext);
				//Sleep(1000); // Not for sure
			}
			continue;
		}

		switch(pIoContext->m_OpType)
		{
			case ACCEPT_POSTED:
					this->DoAccept(pSocketContext,pIoContext);
					break;
			case RECV_POSTED:
					this->DoReceive(pSocketContext,pIoContext);
					break;
			case SEND_POSTED:
					// If the parameter indicate that the buffer need to be release, then release buffer from memory pool
					if(true == pIoContext->m_bNeedRelease)
					{
						MemoryPool::Instance()->Release(pIoContext->m_wsaBuf.buf);
						m_pcClientList->SetTransBytes(pIoContext->m_sockAccept,pIoContext->m_wsaBuf.len,0);
						DateTime clsDateTime;
						clsDateTime.SetToCurrentDateTime();
						m_pcClientList->SetTime(pIoContext->m_sockAccept,clsDateTime);
						RELEASE(pIoContext);
					}
					break;
			case AUTH_POSTED:
				{
					MemoryPool::Instance()->Release(pIoContext->m_wsaBuf.buf);
					RELEASE(pIoContext);
					break;
				}
			case TEST_POSTED:
				{
					m_pcClientList->SetTransBytes(pIoContext->m_sockAccept,pIoContext->m_wsaBuf.len,0);
					DateTime clsDateTime;
					clsDateTime.SetToCurrentDateTime();
					m_pcClientList->SetTime(pIoContext->m_sockAccept,clsDateTime);
					free(pIoContext->m_wsaBuf.buf);
					RELEASE(pIoContext)
					break;
				}
			default:
				break;
		}
	}
}

bool ThreadIOCP::OnNotify(SOCKET p_stSocket, char* p_szBuffer, unsigned int p_uiLength, unsigned int p_uiIndex)
{
	IOCompletionPort* pInstance = IOCompletionPort::Instance();

	// Verify parameters
	if(NULL == p_szBuffer)
	{
		LOG(ERROR) << "Parameter p_szBuffer is NULL";
		return false;
	}
	if(0 == p_uiLength)
	{
		LOG(ERROR) << "Parameter p_iLength is 0";
		return false;
	}

	// WSASend successful message
	char* pszFeedback = this->GetBufferFromMemoryPool(p_uiLength, p_uiIndex);
	if(NULL == pszFeedback)
	{
		MemoryPool* pcMemoryPool = MemoryPool::Instance();
		pcMemoryPool->PrintMemoryStatus();
		LOG(ERROR) << "Cannot allocate memory from memory pool";
		return false;
	}
					
	// Copy message to buffer
	memcpy(pszFeedback,p_szBuffer,p_uiLength);
	PER_IO_CONTEXT* pstSendContext = this->GetNewIOConetext(pszFeedback,p_uiLength,SEND_POSTED,p_stSocket,true);

	// Send content to client
	if(false == pInstance->PostSend(pstSendContext))
	{
		unsigned char ucFileType = *pszFeedback;
		if ((ucFileType&0x80) == 0x80)
		{
			ucFileType = ucFileType & 0x0f;
		}

		LOG(ERROR) << "Post send error. Socket : " << pstSendContext->m_sockAccept << " File Type : " << (unsigned int)(ucFileType);
		this->RemoveClient(pstSendContext);
		return false;
	}

	 return true;
}

void ThreadIOCP::RemoveFromQueue(SOCKET p_iClientSocket)
{
	// Remove client in circle queue
	map<int,CircleDequeueBase*>::iterator it = m_mapCircleQueue.begin();
	while(it != m_mapCircleQueue.end())
	{
		it->second->DeleteSocket(p_iClientSocket);
		it++;
	}
}

void ThreadIOCP::RemoveClient(PER_IO_CONTEXT* p_pstIoContext)
{
	// Remove client from the socket context list
	IOCompletionPort::Instance()->RemoveContext(p_pstIoContext->m_sockAccept);

	// Remove client from client list
	m_pcClientList->DeleteClient(p_pstIoContext->m_sockAccept);

	// Recycle memory
	if (true == p_pstIoContext->m_bNeedRelease && p_pstIoContext->m_OpType == RECV_POSTED && ((p_pstIoContext->m_wsaBuf.len == sizeof(DataHead)) || (p_pstIoContext->m_wsaBuf.len == sizeof(FileType))))
	{
		// Free beat heart packet buffer or receive file request packet buffer
		free(p_pstIoContext->m_wsaBuf.buf);
	}
	else
	{
		// Decrease reference count
		MemoryPool::Instance()->Release(p_pstIoContext->m_wsaBuf.buf);
	}

	// Close socket
	RELEASE_SOCKET(p_pstIoContext->m_sockAccept)

	// Release IO context
	RELEASE(p_pstIoContext);
}